# server_main.py
import numpy as np
import time
import struct
import pickle
import socket
from environment import rm_simple_simulator
from server_client import ServerComm


def main():
    print("Starting server...")
    # Role allocation matrix: rows = [infantry, hero, sentry], cols = [red, blue]
    # 1 = connect client, 0 = random actions
    matrix = [[0, 0], [0, 1], [0, 0]]  # infantry red/blue  # hero  # sentry random

    env = rm_simple_simulator(render_mode="human")  # Enable render for visualization
    print("Environment initialized.")

    print("Initializing server...")
    server = ServerComm()
    server.set_roles(matrix)
    server.accept_connections()

    # Map roles to robot indices (after shuffle)
    print("Mapping roles to robot indices...")
    role_to_idx = {}
    for i, robot in enumerate(env.robots):
        role = (robot["team"], robot["type"])
        if role in server.clients:
            role_to_idx[role] = i

    num_robots = len(env.robots)
    connected_idxs = set(role_to_idx.values())
    dt = 1 / 60.0
    done = False

    # Main loop

    obs, _ = env.reset()
    # send observations at the start
    print("Sending initial observations to clients...")
    for role in list(server.clients.keys()):
        i = role_to_idx[role]
        vision = obs["visions"][i]
        reward = 0
        done_flag = False
        infos = obs["infos"]
        hp_all_dict = {f"{r['team']}_{r['type']}": r["hp"] for r in env.robots}
        info = {
            "remaining_time": env.loader.time - infos["time"],
            "red_score": infos["scores"]["red"],
            "blue_score": infos["scores"]["blue"],
            "hp_all": hp_all_dict,
            "own_hp": infos["hp"][i],
            "own_max_hp": env.robots[i]["max_hp"],
            "bullet": infos["bullet"][i],
            "max_bullet": env.robots[i]["max_bullet"],
            "attack_cd": infos["attack_cooldown"][i],
            "reborn_cd": (
                0
                if env.robots[i]["alive"]
                else env.robots[i]["reborn_time"] - env.robots[i]["reborn_timer"]
            ),
        }
        server.send_observation(role, vision, reward, done_flag, info)

    print(
        f"The size of observation imagee is {vision.shape}, and the sended data size is {len(pickle.dumps(vision))} bytes."
    )

    while not done:
        actions = np.zeros(3 * num_robots, dtype=np.float32)

        # Collect actions from connected clients
        print("Collecting actions from connected clients...")
        for role in list(server.clients.keys()):
            action = server.get_action(role)
            if action is not None:
                i = role_to_idx[role]
                actions[3 * i : 3 * i + 3] = action
            else:
                # Default to zero if failed
                i = role_to_idx[role]
                actions[3 * i : 3 * i + 3] = 0.0

        # Random actions for unconnected robots
        print("Generating random actions for unconnected robots...")
        for i in range(num_robots):
            if i not in connected_idxs:
                rand_slice = env.action_space.sample()[3 * i : 3 * i + 3]
                actions[3 * i : 3 * i + 3] = rand_slice

        # Step environment
        print("Stepping environment...")
        obs, rewards, dones, truncated, infos = env.step(actions)
        done = dones.any() or truncated
        # env.render()

        # Send observations to clients
        print("Sending observations to clients...")
        for role in list(server.clients.keys()):
            i = role_to_idx[role]
            vision = obs["visions"][i]
            reward = rewards[i]
            done_flag = dones[i]
            hp_all_dict = {f"{r['team']}_{r['type']}": r["hp"] for r in env.robots}
            info = {
                "remaining_time": env.loader.time - infos["time"],
                "red_score": infos["scores"]["red"],
                "blue_score": infos["scores"]["blue"],
                "hp_all": hp_all_dict,
                "own_hp": infos["hp"][i],
                "own_max_hp": env.robots[i]["max_hp"],
                "bullet": infos["bullet"][i],
                "max_bullet": env.robots[i]["max_bullet"],
                "attack_cd": infos["attack_cooldown"][i],
                "reborn_cd": (
                    0
                    if env.robots[i]["alive"]
                    else env.robots[i]["reborn_time"] - env.robots[i]["reborn_timer"]
                ),
            }
            server.send_observation(role, vision, reward, done_flag, info)

        time.sleep(dt)

    env.close()
    server.close()


if __name__ == "__main__":
    main()
